// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#pragma once

#include <deque>
#include <limits>
#include <memory>
#include <random>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include <boost/functional/hash.hpp>
#include <glog/logging.h>

#include "coin2/base/request/i_session_pool.h"
#include "coin2/base/request/session.h"
#include "coin2/base/request/util.h"

namespace coin2::base::request {
namespace detail {

struct SessionRtt {
  int64_t rtt;
  std::weak_ptr<Session> session;

  SessionRtt(int64_t r, std::shared_ptr<Session> s) : rtt{r}, session{s} {}
};

bool operator<(const SessionRtt& lhs, const SessionRtt& rhs);

class SingleHostSessionPool : public Session::Callback {
 public:
  SingleHostSessionPool(
      boost::asio::io_context* io_ctx,
      ssl::context* ssl_ctx,
      const std::string& host,
      int port,
      const std::string& bind_ip,
      int max_sessions);  // max_sessions = 0 means no limit on session number.
  ~SingleHostSessionPool();

  void AsyncRequest(
      Session::RequestType req,
      Session::CallbackType callback,
      RequestPriorityWeight priority_weight = RequestPriorityWeight::kMedium);
  void AsyncShutdown();

  void OnTimer(int64_t timeout);
  void ProcessTimeout(int64_t timeout);
  void ProcessOverage(int64_t overage);
  int TotalNumReqs();
  bool is_blocking() const;
  void SetHttpProxy(const std::string& http_proxy) { http_proxy_ = http_proxy; }

 public:
  RttInfo GetRttInfo();
  RequestStatus GetRequestInfo();

 private:
  void OnSessionReady(
      std::weak_ptr<Session> session,
      int64_t rtt,
      const std::string& method,
      const std::string& target,
      int32_t http_code) override;
  void OnSessionError(std::weak_ptr<Session> session) override;

  // for session manage
  auto TakeSession(RequestPriorityWeight priority_weight) -> std::shared_ptr<Session>;
  void DeferDeleteSession(std::weak_ptr<Session> session, bool request_window_need_updating);
  void DeleteSession();
  // for statistic
  void StatsRequestInfo(const std::string& req_method, const std::string& req_target, int32_t code);
  // for handling timed out session
  void ShutdownTimeoutSession(int64_t now);
  // for choosing fastest session
  auto FindMinOrMaxRttSession(bool isMin) -> std::unordered_set<std::shared_ptr<Session>>::iterator;
  // for session assigning
  void UpdateRequestWindow(int64_t timestamp, RequestPriorityWeight priority_weight);
  bool CheckIfFreeSessionsAreEnoughForQuery(RequestPriorityWeight priority_weight);
  bool WhetherNeedToPreCreateSession();
  int PoissonDistributionSampling();

 private:
  boost::asio::io_context* io_ctx_;
  ssl::context* ssl_ctx_;
  const std::string host_;
  const int port_;
  const std::string bind_ip_;
  std::string http_proxy_;
  std::shared_ptr<detail::SimpleObjectPool<Session::RequestContext>> req_ctx_pool_;

  // for session manage
  int max_sessions_ = 0;
  std::unordered_set<std::shared_ptr<Session>> work_sessions_;
  std::unordered_set<std::shared_ptr<Session>> free_sessions_;
  std::unordered_set<std::shared_ptr<Session>> del_sessions_;
  std::unordered_set<std::shared_ptr<Session>> timeout_sessions_;

  // for statistic
  RequestStatus req_status_;
  RttInfo rtts_;
  bool is_prex_{false};

  // for handling timed out session
  int64_t last_shutdown_timeout_session_ts_ = 0;
  const static int64_t shutdown_timeout_session_ns_ = 5 * 60 * 1000'000'000L;
  const static int64_t check_shutdown_timeout_session_period_ns_ = 60 * 1000'000'000L;

  // for session assigning
  int session_window_size_ = 10;
  int working_query_num_ = 0;
  double min_proportion_of_query_ = 0.2;
  std::deque<int64_t> free_session_add_records_list_;
  std::deque<int64_t> free_session_del_records_list_;
  std::default_random_engine generator_;

  // for choosing fastest session
  std::unordered_map<std::shared_ptr<Session>, std::vector<int64_t>> sessions_rtt_list_;
  std::set<SessionRtt> rtt_tree_;
};

}  // namespace detail

class SessionPool : public ISessionPool, public std::enable_shared_from_this<SessionPool> {
 public:
  explicit SessionPool(boost::asio::io_context* io_ctx);
  SessionPool(boost::asio::io_context* io_ctx, int max_sessions_per_host);
  SessionPool(boost::asio::io_context* io_ctx, ssl::context* ssl_ctx, int max_sessions_per_host);
  ~SessionPool();

  void AsyncRequest(
      const std::string& host,
      int port,
      const std::string& bind_ip,
      Session::RequestType req,
      Session::CallbackType callback,
      RequestPriorityWeight priority_weight = RequestPriorityWeight::kMedium,
      http::verb method = http::verb::get);
  void AsyncRequest(
      const std::string& host,
      int port,
      const std::string& bind_ip,
      std::string_view target,
      Session::CallbackType callback,
      http::verb method = http::verb::get);
  void AsyncRequest(
      const std::string& host,
      int port,
      Session::RequestType req,
      Session::CallbackType callback,
      http::verb method = http::verb::get) {
    AsyncRequest(host, port, "", std::move(req), callback, RequestPriorityWeight::kMedium, method);
  }
  void AsyncRequest(
      const std::string& host,
      int port,
      std::string_view target,
      Session::CallbackType callback,
      http::verb method = http::verb::get) {
    AsyncRequest(host, port, "", target, callback, method);
  }

  // inherit from ISessionPool
  void StartTimer() override;
  void StopTimer() override { timer_enabled_ = false; };
  void SetTimeout(int64_t timeout_ns) override;
  bool is_blocking() const override;
  void AsyncRequest(
      const std::string& host,
      int port,
      const std::string& bind_ip,
      Session::RequestType req,
      Session::CallbackType callback,
      RequestPriorityWeight priority_weight) override;
  void SetHttpProxy(const std::string& http_proxy) override { http_proxy_ = http_proxy; };
  RttInfo GetRttInfo() override;
  RequestStatus GetRequestInfo() override;

  void AsyncShutdown();

  template <typename Arg, typename CompletionToken>
  auto CoroRequest(
      const std::string& host,
      int port,
      Arg&& arg,
      CompletionToken&& token,
      http::verb method = http::verb::get) {
    using Result = boost::asio::
        async_result<CompletionToken, void(boost::system::error_code, const SessionResponse*)>;
    using Handler = typename Result::completion_handler_type;

    Handler handler(std::forward<CompletionToken>(token));
    Result result(handler);
    AsyncRequest(host, port, std::forward<Arg>(arg), handler, method);
    return result.get();
  }

 private:
  using Key = std::tuple<std::string, int, std::string>;
  void OnTimer(const boost::system::error_code& ec);
  detail::SingleHostSessionPool*
  GetSingleHostSessionPool(const std::string& host, int port, const std::string& bind_ip);

  std::unordered_map<Key, std::unique_ptr<detail::SingleHostSessionPool>, boost::hash<Key>> pools_;
  boost::asio::io_context* io_ctx_;
  ssl::context* ssl_ctx_;
  int max_sessions_per_host_;

  bool timer_enabled_ = false;
  boost::asio::steady_timer timer_;
  int64_t timeout_ns_ = 30 * 1000000000L;
  int64_t overage_ns_ = std::numeric_limits<int64_t>::max();

  std::string http_proxy_ = "";
};

std::string GetHttpsContent(const std::string& url_root, const std::string& url_path);
}  // namespace coin2::base::request
